import torch
import torch.nn as nn

class CNN_BE(nn.Module):
    def __init__(self, in_c, out_c=1, mid_planes=64):
        super().__init__()
        self.mid_planes = mid_planes
        self.layers = 7
        self.conv0 = nn.Conv2d(in_c, mid_planes, kernel_size=3, stride=1, padding=1,padding_mode='circular', bias=False)
        self.conv1 = nn.Conv2d(mid_planes,mid_planes,kernel_size=3, stride=1, padding=1, padding_mode='circular',bias=False)
        # self.bn1 = nn.BatchNorm2d(mid_planes)
        self.conv2 = nn.Conv2d(mid_planes,mid_planes,kernel_size=3, stride=1, padding=1, padding_mode='circular',bias=False)
        # self.bn2 = nn.BatchNorm2d(mid_planes)
        self.conv3 = nn.Conv2d(mid_planes, mid_planes, kernel_size=3, stride=1, padding=1, padding_mode='circular',bias=False)
        self.conv4 = nn.Conv2d(mid_planes, mid_planes, kernel_size=3, stride=1, padding=1, padding_mode='circular',bias=False)
        self.conv5 = nn.Conv2d(mid_planes, mid_planes, kernel_size=3, stride=1, padding=1, padding_mode='circular',bias=False)
        self.conv6 = nn.Conv2d(mid_planes, mid_planes, kernel_size=3, stride=1, padding=1, padding_mode='circular',bias=False)
        self.conv7 = nn.Conv2d(mid_planes, mid_planes, kernel_size=3, stride=1, padding=1, padding_mode='circular',
                               bias=False)
        self.conv8 = nn.Conv2d(mid_planes, mid_planes, kernel_size=3, stride=1, padding=1, padding_mode='circular',
                               bias=False)
        self.conv9 = nn.Conv2d(mid_planes, mid_planes, kernel_size=3, stride=1, padding=1, padding_mode='circular',
                               bias=False)
        self.conv10 = nn.Conv2d(mid_planes, mid_planes, kernel_size=3, stride=1, padding=1, padding_mode='circular',
                               bias=False)
        self.conv11 = nn.Conv2d(mid_planes, mid_planes, kernel_size=3, stride=1, padding=1, padding_mode='circular',
                               bias=False)
        self.conv12 = nn.Conv2d(mid_planes, mid_planes, kernel_size=3, stride=1, padding=1, padding_mode='circular',
                               bias=False)
        self.conv13 = nn.Conv2d(mid_planes, mid_planes, kernel_size=3, stride=1, padding=1, padding_mode='circular',
                               bias=False)
        self.conv14 = nn.Conv2d(mid_planes, mid_planes, kernel_size=3, stride=1, padding=1, padding_mode='circular',
                                bias=False)
        self.convF = nn.Conv2d(mid_planes, out_c,kernel_size=3, stride=1, padding=0)

    def forward(self,x):
        residual = self.conv0(x)

        out = self.conv0(x)

        #resblock1
        out = self.conv1(out)
        out = torch.tanh(out)
        out = self.conv2(out)
        out = residual+out
        out = torch.tanh(out)
        #resblock2
        residual = out
        out = self.conv3(out)
        out = torch.tanh(out)
        out = self.conv4(out)
        out = residual + out
        out = torch.tanh(out)

        #resblock3
        residual = out
        out = self.conv5(out)
        out = torch.tanh(out)
        out = self.conv6(out)
        out = residual + out
        out = torch.tanh(out)
        #
        # # # resblock4
        residual = out
        out = self.conv7(out)
        out = torch.tanh(out)
        out = self.conv8(out)
        out = residual + out
        out = torch.tanh(out)
        #
        # # resblock5
        residual = out
        out = self.conv9(out)
        out = torch.tanh(out)
        out = self.conv10(out)
        out = residual + out
        out = torch.tanh(out)
        
        # # resblock6
        # residual = out
        out = self.conv11(out)
        out = torch.tanh(out)
        out = self.conv12(out)
        out = residual + out
        out = torch.tanh(out)
        #
        # # resblock7
        residual = out
        out = self.conv13(out)
        out = torch.tanh(out)
        out = self.conv14(out)
        out = residual + out
        out = torch.tanh(out)
        out = torch.tanh(out)
        out = self.convF(out)

        return out
    
    def name(self):
        return f'cnnBE-{self.mid_planes}-{self.layers}'